import ctypes
from ctypes import *
import numpy as np
import time
from numpy.ctypeslib import ndpointer

dll = ctypes.cdll.LoadLibrary
lib = dll("./ascend_index.so")

# Params
device_id = 1
dim = 512 
n = 1000000
length = n * dim
topK = 10
search_num = 48
q_length_arr = search_num * dim
q_length = topK * search_num
loop_times = 100

# init index
lib.init(device_id, dim)

# gen feature
print("============ gen feature =================")
start = time.time()
r = np.random.rand(length).astype(np.float32)
end = time.time()
print(f"========= gen feature cost: {end - start} =======")

# insert feature
print("=========== insert feature ================")
start = time.time()
insert_gallery = lib.insert_gallery
insert_gallery.argtypes = [c_int, ndpointer(ctypes.c_float)]
insert_gallery.restype = None
insert_gallery(n, r)
end = time.time()
print(f"============ insert feature cost: {end - start}")

# get baseSize
base_size = lib.getBaseSize()
print("baseSize: ", base_size)

# search
print("=============== search ================")
search = lib.search
search.argtypes = [c_int, ndpointer(ctypes.c_float), c_int, ndpointer(ctypes.c_float), ndpointer(ctypes.c_int64)]
search.restype = None
q_test = np.random.rand(q_length_arr).astype(np.float32)
scores_c = np.zeros(q_length).astype(np.float32)
idxs_c = np.zeros(q_length).astype(np.int64)
start = time.time()
for _ in range(loop_times):
    search(search_num, q_test, topK, scores_c, idxs_c)
end = time.time()
print(f"============= {end - start} ==============")
print(f"Flat dim={dim}, base={n}, batch={search_num}, topK={topK}, qps={(loop_times * search_num) / (end - start)}")








